import codecs
import json

from gxl_ai_utils.utils import utils_file

def do_handle():
    """
    转移数据， 熟悉数据，可能的数据翻译
    :return:
    """
    # 数据翻译
    from pygtrans import Translate

    # client = Translate(proxies={'https': 'http://localhost:7890'})
    client = Translate()
    # 检测语言
    # text = client.detect('Answer the question.')
    # assert text.language == 'en'
    # 翻译句子
    text = client.translate('Look at these pictures and answer the questions.')
    print(text.translatedText)


def do_handle_2():
    """
    对目录进行打tar包
    :return:
    """
    input_dir="/home/node54_tmpdata/xlgeng/caption_data/shards"
    output_dir = "/home/node54_tmpdata/xlgeng/caption_data/shards_2"
    num = 220
    for i in range(num):
        print(f'进度： {i} / {num}')
        dir_name = str(i).zfill(3)
        extract_path = f"{input_dir}/Audio_Caption_000000{dir_name}"
        utils_file.do_compress_directory_by_tar_form(extract_path, output_dir)

def do_handle_3():
    """
    上述的用python对目录打成tar包的文件程序都读取失败了。 先转成raw格式，然后使用标准脚本进行打包
    :return:
    """
    res_dict_list = utils_file.load_dict_list_from_jsonl('./data/data.list')
    for i in range(558, 634):
        tmp_data_list = []
        i_str = str(i).zfill(3)
        input_dir = f"/home/node54_tmpdata/xlgeng/caption_data/shards/Audio_Caption_000000{i_str}"
        print(f'input_dir : {input_dir}')
        # task,key, wav,txt,lang, speaker, emotion,gender,duration,extra{duration:}
        key_list = ['task', 'wav', 'txt', 'lang', 'speaker', 'emotion', 'gender', 'duration', 'extra']
        for tmp_key in key_list:
            tmp_dict = utils_file.do_get_scp_for_wav_dir(input_dir, suffix=tmp_key)
            new_tmp_dict = {}
            if tmp_key == 'wav':
                for k, v in tmp_dict.items():
                    value_new = v
                    new_tmp_dict[k] = value_new
            elif tmp_key == 'extra':
                for k, v in tmp_dict.items():
                    with codecs.open(v, 'r', encoding='utf=8') as f:
                        value_new = json.load(f)
                    new_tmp_dict[k] = value_new
            elif tmp_key == 'lang':
                for k, v in tmp_dict.items():
                    value_new = "<CN>"
                    new_tmp_dict[k] = value_new
            else:
                for k, v in tmp_dict.items():
                    value_new = utils_file.load_first_row_clean(v)
                    new_tmp_dict[k] = value_new
            tmp_dict = new_tmp_dict
            # print(tmp_dict)
            if len(tmp_data_list) == 0:
                for k, v in tmp_dict.items():
                    added_dict = {tmp_key: v, 'key': k}
                    tmp_data_list.append(added_dict)
            else:
                assert len(tmp_data_list) == len(tmp_dict)
                for res_dict in tmp_data_list:
                    key_now = res_dict['key']
                    if key_now in tmp_dict:
                        res_dict[tmp_key] = tmp_dict[key_now]
        res_dict_list.extend(tmp_data_list)
        if i % 10 == 0:
            output_jsonl_path = "./data/data_gxl2.list"
            utils_file.write_dict_list_to_jsonl(res_dict_list, output_jsonl_path)
    output_jsonl_path = "./data/data_gxl2.list"
    utils_file.write_dict_list_to_jsonl(res_dict_list, output_jsonl_path)





if __name__ == '__main__':
    do_handle_3()
